import glob
import os
import shutil
import matplotlib.pyplot as plt
import numpy as np
import tqdm
import sys
sys.path.append(os.path.dirname(os.path.realpath(__file__)))
from punctuation import process_text
input_data_dir = '/home/work_nfs5_ssd/hfxue/data/data4w/source_1/huiting400'
output_data_dir = '/home/backup_nfs5/data_tts/wenetspeech'

py_source = '/home/backup_nfs5/data_tts/wenetspeech'
from gxl_ai_utils.utils import utils_file
from gxl_ai_utils.AiConstant import AI_LOGGER
logger = AI_LOGGER('./output/log/asr_handle_pipline.log')

def make_scp(root_dir:str, output_dir:str, file_name:str='wav.scp'):
    utils_file.makedir_sil(output_dir)
    res_dic = {}
    for root, dirs, files in os.walk(root_dir):
        for dirname in dirs:
            """"""
            dir_path = os.path.join(root, dirname)
            files = glob.glob(f"{dir_path}/*.wav")
            for file_path in tqdm(files, total=len(files), desc=dirname):
                """"""
                filename = file_path.split("/")[-1]
                key = filename.split('.')[0]
                res_dic[key] = file_path
        break
    utils_file.write_dic_to_scp(res_dic, os.path.join(output_dir,file_name))
    print('over')

def make_mos_scp(root_dir:str, output_dir:str, file_name:str='mos.scp'):
    utils_file.makedir_sil(output_dir)
    res_dict = {}
    for root, dirs, files in os.walk(root_dir):
        for file in files:
            # 检查文件扩展名是否为.scp
            if file.endswith(".scp"):
                # 构建完整的文件路径
                file_path = os.path.join(root, file)
                temp_dict = utils_file.load_dic_from_scp(file_path)
                res_dict.update(temp_dict)
    utils_file.write_dic_to_scp(res_dict, os.path.join(output_dir,file_name))


def show_image(dataset_name: str, input_list: list, output_dir: str):
    # 生成包含10000个数值的列表，范围为0-5之间
    utils_file.makedir_sil(output_dir)
    data = input_list
    # 设置梯度为0.5
    bins = np.arange(0, 5.1, 0.5)
    # 绘制直方图
    plt.hist(data, bins=bins, edgecolor='black', alpha=0.7)

    # 标明每个柱子的具体数量
    for i in range(len(bins) - 1):
        count = np.sum((data >= bins[i]) & (data < bins[i + 1]))
        plt.text(bins[i] + 0.25, count + 50, str(count), ha='center', va='bottom')

    # 设置图表标题和标签
    plt.title('Histogram of {}'.format(dataset_name))
    plt.xlabel('Value')
    plt.ylabel('Frequency')

    # 显示图表
    plt.savefig(os.path.join(output_dir, 'histogram_{}.png'.format(dataset_name)))

if __name__ == '__main__':
    dataset_name = 'wenetspeech'
    logger.info('开始得到mos.scp')
    make_mos_scp(os.path.join(output_data_dir, 'mos'), output_data_dir, 'mos.scp')
    logger.info('mos.scp得到完毕')

    logger.info('开始绘制直方图')
    mos_dict = utils_file.load_dic_from_scp(os.path.join(output_data_dir, 'mos.scp'))
    mos_num_list = list(mos_dict.values())
    float_array = np.array(mos_num_list, dtype=float)
    show_image(dataset_name, float_array, output_data_dir)
    logger.info('直方图绘制完毕')